'''
Description: 
Author: 
LastEditors: Please set LastEditors
Date: 2022-12-08 15:06:11
LastEditTime: 2022-12-08 15:09:17
'''
import sys
import os
import time
import cv2
from loguru import logger

sys.path.append("D:\AI\YOLOX")  #YOLOX 位置

file_idx = 0
file_save_path = 'd:\\AI\\YOLOX\\datasets\\coco\\unlabel' #存储位置
save_interval = 1
time_temp = time.time()


def __saveImage(frame):
    global file_idx
    global file_save_path
    global time_temp

    t0 = time.time()
    if time.time()-time_temp < save_interval:
        return None
    time_temp = t0

    file_idx += 1
    file_name = str(file_idx).zfill(6)

    # 储存jpg图片
    jpg_name = file_name+'.jpg'
    jpg_path = os.path.join(file_save_path, jpg_name)
    logger.debug("save image : {}".format(jpg_path))
    cv2.imwrite(jpg_path, frame)
   
    return jpg_path


def __getMaxFileNameId(filePath: str = None, fileMaxId: int = 0):
    assert filePath is not None, "file path is None"
    fileList = os.listdir(filePath)
    maxId = fileMaxId
    for fileWithPath in fileList:
        if not fileWithPath.endswith(".jpg"):
            continue
        # dir, fileName = os.path.split(fileWithPath)  # 分离路径和文件名
        name, ext = os.path.splitext(fileWithPath)  # 分离文件名和扩展名
        if not name.isdigit():  # 判断是否是数字
            continue
        n = int(name)
        if n > maxId:
            maxId = n
    return maxId


if __name__ == "__main__":

    # 每隔一段时间存储一张图像，用于训练
    save_interval = 1
    os.makedirs(file_save_path, exist_ok=True)
    file_idx = __getMaxFileNameId(file_save_path)  # 文件名编号

    camera_id = 0
    cap = cv2.VideoCapture(camera_id, cv2.CAP_DSHOW)

    if not cap.isOpened():
        logger.error("can not open camera {}".format(camera_id))

    cap.set(cv2.CAP_PROP_FRAME_WIDTH,800) #设置摄像头分辨率
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT,600)
    

    while cap.isOpened():
        success, frame = cap.read()
        assert success, 'no frame read'     
        
        cv2.imshow('my_window', frame)    

        __saveImage(frame)    

        ch = cv2.waitKey(1)
        if ch == 27 or ch == ord("q") or ch == ord("Q"):
            break
    cap.release()

    cv2.destroyAllWindows()
